package com.lowcode.service.impl;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.benjaminwan.ocrlibrary.OcrResult;
import com.lowcode.entity.DaLabel;
import com.lowcode.entity.DaLabelImages;
import com.lowcode.service.DaLabelImagesService;
import com.lowcode.service.DaLabelService;
import com.lowcode.mapper.DaLabelMapper;
import com.lowcode.vo.Result;
import io.github.mymonstercat.Model;
import io.github.mymonstercat.ocr.InferenceEngine;
import io.github.mymonstercat.ocr.config.ParamConfig;
import lombok.extern.slf4j.Slf4j;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.PDFRenderer;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;

/**
* @author guowenhao
* @description 针对表【da_label】的数据库操作Service实现
* @createDate 2024-03-18 21:03:58
*/
@Slf4j
@Service
public class DaLabelServiceImpl extends ServiceImpl<DaLabelMapper, DaLabel>
    implements DaLabelService{

    @Autowired
    private DaLabelImagesService daLabelImagesService;

    @Override
    @Scheduled(cron = "0/30 * * * * ?")
    public Result<?> ocr() {
        // 获取一条数据
        DaLabel daLabel = this.getOne(new LambdaQueryWrapper<DaLabel>().eq(DaLabel::getStatus,0).last("LIMIT 1"));

        if (StringUtils.isEmpty(daLabel.getTitleUrl())) {
            // 修改label的status
            daLabel.setStatus(1);
            this.updateById(daLabel);
            return Result.error("PDF地址为空");
        }

        // 将PDF转换成PNG图片
        List<String> pngList;
        try {
            // 将PDF路径切片，得到pdf路径
            String[] parts = daLabel.getTitleUrl().split(daLabel.getRecordNumber());
            pngList = this.convertPdfToImage(daLabel, parts[0], daLabel.getRecordNumber());
        } catch (IOException e) {
            e.printStackTrace();
            return Result.error(e.getMessage());
        }

        // 同示例中的配置
        ParamConfig paramConfig = ParamConfig.getDefaultConfig();
        paramConfig.setDoAngle(true);
        paramConfig.setMostAngle(true);
        InferenceEngine engine = InferenceEngine.getInstance(Model.ONNX_PPOCR_V3);

        List<DaLabelImages> daLabelImagesList = new ArrayList<DaLabelImages>();

        for (String url : pngList) {
            // 开始识别
            OcrResult ocrResult = engine.runOcr(url, paramConfig);
            // 根据识别结果，导入label_images
            DaLabelImages daLabelImages = new DaLabelImages();
            daLabelImages.setLabelId(daLabel.getId());
            daLabelImages.setImage(url);
            daLabelImages.setPngUrl(url);
            daLabelImages.setContentText(ocrResult.getStrRes().toString());
            System.out.println("识别完成：" + url);
            daLabelImagesList.add(daLabelImages);
        }
        // 插入数据
        boolean res = daLabelImagesService.saveOrUpdateBatch(daLabelImagesList);
        if (res) {
            // 修改label的status
            daLabel.setStatus(1);
            this.updateById(daLabel);
        }
        return Result.OK();
    }


    public List<String> convertPdfToImage(DaLabel daLabel, String outputDir, String pdfName) throws IOException {
        List<String> pngList = new ArrayList<>();
        // 创建输出目录
        Path outputDirPath = Paths.get(outputDir);
        if (!Files.exists(outputDirPath)) {
            Files.createDirectories(outputDirPath);
        }

        try (PDDocument document = PDDocument.load(new File(daLabel.getTitleUrl()))) {
            PDFRenderer pdfRenderer = new PDFRenderer(document);
            for (int page = 0; page < document.getNumberOfPages(); ++page) {
                BufferedImage bim = pdfRenderer.renderImageWithDPI(page, 300); // DPI设置为300，可以根据需要调整
                String fileName = outputDir + "/" + pdfName + "_page_" + (page + 1) + ".png";
                ImageIO.write(bim, "png", new File(fileName));
                pngList.add(fileName);
            }
        } catch (Exception e) {
            log.info("未找到文件：" + daLabel.getTitleUrl());
            daLabel.setStatus(1);
            this.updateById(daLabel);
        }

        return pngList;
    }

}




